package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"log"
	"os"
	"strconv"
	"strings"
)

func getObj(results []entity.Column, name string) entity.Column {
	var obj entity.Column
	for _, result := range results {
		if result.Name() == name {
			obj = result
		}
	}
	return obj
}

// 辅助函数：将 float32 数组转换为逗号分隔的字符串
func float32ArrayToString(array []float32) string {
	var strArray []string
	for _, val := range array {
		strArray = append(strArray, strconv.FormatFloat(float64(val), 'f', -1, 32))
	}
	return strings.Join(strArray, ",")
}

func main() {
	// 创建 Milvus 客户端
	ctx := context.Background()
	milvusClient, err := client.NewGrpcClient(ctx, "192.168.30.80:31530") // 修改为你的 Milvus 地址
	if err != nil {
		log.Fatalf("failed to create milvus client: %v", err)
	}
	defer milvusClient.Close()

	// 查询条件和需要输出的字段
	collectionName := "galleryNew"
	outputFields := []string{
		"embedding", "image_type", "gallery_image_id", "path",
		"ai_model", "model_version", "prompt", "width", "height",
		"aspectRatio", "created_at", "deleted_at",
	}

	var totalRecords = 130000
	var batchSize = 1000
	var fileBatchSize = 2000
	var totalFiles = totalRecords / fileBatchSize

	// 记录上次查询的最后一个 `created_at`，初始为 0
	var lastCreatedAt int64 = 0

	// 分批获取并写入数据
	for fileNum := 0; fileNum < totalFiles; fileNum++ {
		// 创建 CSV 文件，文件名递增
		fileName := fmt.Sprintf("milvus_query_results_%d.csv", fileNum+1)
		file, err := os.Create(fileName)
		if err != nil {
			log.Fatalf("failed to create CSV file: %v", err)
		}
		defer file.Close()

		// 初始化 CSV 写入器
		writer := csv.NewWriter(file)
		defer writer.Flush()

		// 写入表头
		writer.Write(outputFields)

		for batch := 0; batch < fileBatchSize/batchSize; batch++ {
			//offset := (fileNum*fileBatchSize + batch*batchSize)
			opt := client.SearchQueryOptionFunc(func(option *client.SearchQueryOption) {
				option.Limit = int64(batchSize)
				option.Offset = 0
			})

			// 添加 `created_at` 的过滤条件，确保每次查询都是从上次的最后一个时间戳继续
			expr := fmt.Sprintf("created_at > %d", lastCreatedAt)

			// 执行查询，并按 `created_at` 升序排序
			queryResults, err := milvusClient.Query(ctx, collectionName, nil, expr, outputFields, opt)
			if err != nil {
				log.Fatalf("failed to query milvus: %v", err)
			}

			// 遍历查询结果并写入 CSV
			var embeddingObj = getObj(queryResults, "embedding")
			var imageTypeObj = getObj(queryResults, "image_type")
			var galleryImageIdObj = getObj(queryResults, "gallery_image_id")
			var pathObj = getObj(queryResults, "path")
			var aiModelObj = getObj(queryResults, "ai_model")
			var modelVersionObj = getObj(queryResults, "model_version")
			var promptObj = getObj(queryResults, "prompt")
			var widthObj = getObj(queryResults, "width")
			var heightObj = getObj(queryResults, "height")
			var aspectRatioObj = getObj(queryResults, "aspectRatio")
			var createdAtObj = getObj(queryResults, "created_at")
			var deletedAtObj = getObj(queryResults, "deleted_at")

			for i := 0; i < batchSize; i++ {
				// 假设 embedding 是一个 []float32 数组
				embedding, _ := embeddingObj.Get(i)
				embeddingArray, ok := embedding.([]float32)
				if !ok {
					log.Fatalf("unexpected type, expected []float32")
				}
				// 将 float32 切片转换为字符串
				embeddingResult := "[" + float32ArrayToString(embeddingArray) + "]"

				imageType, _ := imageTypeObj.GetAsInt64(i)
				galleryImageId, _ := galleryImageIdObj.GetAsInt64(i)
				path, _ := pathObj.GetAsString(i)
				aiModel, _ := aiModelObj.GetAsString(i)
				modelVersion, _ := modelVersionObj.GetAsString(i)
				prompt, _ := promptObj.GetAsString(i)
				width, _ := widthObj.GetAsInt64(i)
				height, _ := heightObj.GetAsInt64(i)
				aspectRatio, _ := aspectRatioObj.GetAsDouble(i)
				createdAt, _ := createdAtObj.GetAsInt64(i)
				deletedAt, _ := deletedAtObj.GetAsInt64(i)

				record := []string{
					fmt.Sprintf("%v", embeddingResult),
					fmt.Sprintf("%v", imageType),
					fmt.Sprintf("%v", galleryImageId),
					fmt.Sprintf("%v", path),
					fmt.Sprintf("%v", aiModel),
					fmt.Sprintf("%v", modelVersion),
					fmt.Sprintf("%v", prompt),
					fmt.Sprintf("%v", width),
					fmt.Sprintf("%v", height),
					fmt.Sprintf("%v", aspectRatio),
					fmt.Sprintf("%v", createdAt),
					fmt.Sprintf("%v", deletedAt),
				}

				// 写入 CSV
				writer.Write(record)

				// 更新最后一个 `created_at` 值
				lastCreatedAt = createdAt
			}

			// 刷新缓冲区，确保批量写入后立即写入文件
			writer.Flush()
		}

		fmt.Printf("CSV 文件 %s 导出成功！\n", fileName)
	}
}
